package com.test.langchain4j.controller;

import com.test.langchain4j.service.ChatAssistant;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 2025/7/31 15:53
 * @modifiedBy:
 * @version: 1.0
 */
@RestController
@Slf4j
public class StreamingChatModelController {
  // 直接使用 low-level LLM API
  @Resource private StreamingChatModel streamingChatLanguageModel;
  // 自己封装接口使用 high-level LLM API
  @Resource private ChatAssistant chatAssistant;

  /**
   * Chat flux. http://localhost:9007/chatstream/chat?prompt=天津有什么好吃的
   *
   * @param prompt the prompt
   * @return the flux
   */
  @GetMapping(value = "/chatstream/chat")
  public Flux<String> chat(@RequestParam("prompt") String prompt) {
    return Flux.create(
        stringFluxSink -> {
          streamingChatLanguageModel.chat(
              prompt,
              new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                  stringFluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse completeResponse) {
                  stringFluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                  stringFluxSink.error(throwable);
                }
              });
        });
  }

  @GetMapping(value = "/chatstream/chat2")
  public void chat2(@RequestParam(value = "prompt", defaultValue = "北京有什么好吃") String prompt) {
    System.out.println("---come in chat2");
    streamingChatLanguageModel.chat(
        prompt,
        new StreamingChatResponseHandler() {
          @Override
          public void onPartialResponse(String partialResponse) {
            System.out.println(partialResponse);
          }

          @Override
          public void onCompleteResponse(ChatResponse completeResponse) {
            System.out.println("---response over: " + completeResponse);
          }

          @Override
          public void onError(Throwable throwable) {
            throwable.printStackTrace();
          }
        });
  }

  @GetMapping(value = "/chatstream/chat3")
  public Flux<String> chat3(
      @RequestParam(value = "prompt", defaultValue = "南京有什么好吃") String prompt) {
    System.out.println("---come in chat3");

    return chatAssistant.chatFlux(prompt);
  }
}
